package org.example.producer;

import org.apache.rocketmq.client.exception.MQClientException;
import org.apache.rocketmq.client.producer.LocalTransactionState;
import org.apache.rocketmq.client.producer.SendResult;
import org.apache.rocketmq.client.producer.TransactionListener;
import org.apache.rocketmq.client.producer.TransactionMQProducer;
import org.apache.rocketmq.common.message.Message;
import org.apache.rocketmq.common.message.MessageExt;
import org.apache.rocketmq.remoting.common.RemotingHelper;

import java.io.IOException;
import java.io.UnsupportedEncodingException;
import java.util.concurrent.*;

/**
 * @author Tang
 * @classname TransactionProducer
 * @description [ 事物消息只和生产者有关和消费者无关 ]
 * 事务消息发送分为两个阶段。第一阶段会发送一个半事务消息，
 * 半事务消息是指暂不能投递的消息，生产者已经成功地将消息发送到了 Broker，但是Broker 未收到生产者对该消息的二次确认，
 * 此时该消息被标记成“暂不能投递”状态，如果发送成功则执行本地事务，
 * 并根据本地事务执行成功与否，向 Broker 半事务消息状态（commit或者rollback），
 * 半事务消息只有 commit 状态才会真正向下游投递。
 * @date 2022/11/27 16:12
 */
public class TransactionProducer {

    public static void main(String[] args) throws MQClientException, InterruptedException {
        TransactionListener transactionListener = new TransactionListenerImpl();
        TransactionMQProducer producer = new TransactionMQProducer("TransactionProducer");
        producer.setNamesrvAddr("localhost:9876");
        ExecutorService executorService = new ThreadPoolExecutor(2, 5, 100, TimeUnit.SECONDS, new ArrayBlockingQueue<Runnable>(2000), new ThreadFactory() {
            @Override
            public Thread newThread(Runnable r) {
                Thread thread = new Thread(r);
                thread.setName("client-transaction-msg-check-thread");
                return thread;
            }
        });

        producer.setExecutorService(executorService);
        producer.setTransactionListener(transactionListener);
        producer.start();

        String[] tags = new String[] {"TagA", "TagB", "TagC", "TagD", "TagE"};
        for (int i = 0; i < 10; i++) {
            try {
                Message msg =
                        new Message("TestTopic", tags[i % tags.length], "KEY" + i,
                                ("Hello RocketMQ " + i).getBytes(RemotingHelper.DEFAULT_CHARSET));
                SendResult sendResult = producer.sendMessageInTransaction(msg,null);
                System.out.printf("%s%n", sendResult);
                Thread.sleep(10);
            } catch (MQClientException | UnsupportedEncodingException e) {
                e.printStackTrace();
            }
        }
        try {
            System.in.read();
        } catch (IOException e) {
            e.printStackTrace();
        }
        producer.shutdown();
    }

    static class TransactionListenerImpl implements TransactionListener {

        // 是半事务消息发送成功后，执行本地事务的方法
        @Override
        public LocalTransactionState executeLocalTransaction(Message msg, Object arg) {
            if ("TagA".equals(msg.getTags())) {
                return LocalTransactionState.COMMIT_MESSAGE;
            }
            if ("TagB".equals(msg.getTags())) {
                return LocalTransactionState.ROLLBACK_MESSAGE;
            }
//            该方法中返回以下三种状态：
//            LocalTransactionState.COMMIT_MESSAGE：提交事务，允许消费者消费该消息
//            LocalTransactionState.ROLLBACK_MESSAGE：回滚事务，消息将被丢弃不允许消费。
//            LocalTransactionState.UNKNOW：暂时无法判断状态，等待固定时间以后Broker端根据回查规则向生产者进行消息回查。
            return LocalTransactionState.UNKNOW;
        }

        // checkLocalTransaction是由于二次确认消息没有收到，Broker端回查事务状态的方法。
        @Override
        public LocalTransactionState checkLocalTransaction(MessageExt msg) {
            if ("TagC".equals(msg.getTags())) {
                return LocalTransactionState.COMMIT_MESSAGE;
            }
            return LocalTransactionState.UNKNOW;
        }
    }
}
